// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: junglekim
#pragma once

#include "coin2/strategy/order_executor/layering_executor3.h"
#include "coin2/strategy/vmm/vmm_dh_util.h"
#include "coin2/strategy/vmm/vmm_pricer.h"
#include "coin2/strategy/vmm/vmm_util.h"

using coin::proto::OrderSide;
using coin2::exchange::base::feed::FeedTimer;
using coin2::strategy::order_executor::LayeringExecutor3;

namespace coin2::strategy::vmm {
class VmmArbitrageur {
 public:
  VmmArbitrageur(
      const ::nlohmann::json& configs,
      const bool& dry_run,
      const std::string& tag,
      const bool& verbose)
      : configs_(configs), dry_run_(dry_run), verbose_(verbose), tag_(tag) {}
  virtual void InitExecutor(
      std::string symbol,
      int64_t start_time,
      PassiveExecutorConfig config,
      IOrderSubsystem* os,
      Driver* driver,
      const std::optional<double>& qty_cancel_pct,
      const bool is_randomised_post,
      const int cancel2order_ratio) {
    os_[symbol] = os;
    driver_ = driver;
    auto arbi_config = new PassiveExecutorConfig();
    arbi_config->MergeFrom(config);
    arbi_config->set_filter_tag(tag_);
    auto get_product_holder = [this](const IProduct& product, int64_t timestamp) {
      return driver_->product_holder_manager()->GetProductHolder(product, timestamp);
    };
    executors_[symbol] = std::make_unique<LayeringExecutor3>(
        *arbi_config,
        os,
        start_time,
        nullptr,
        get_product_holder,
        std::nullopt,
        std::nullopt,
        qty_cancel_pct,
        is_randomised_post,
        cancel2order_ratio,
        verbose_);
  }
  virtual bool onPricerUpdate(
      const std::unordered_map<std::string, VmmPriceResult>& pricer_results,
      const int64_t& timestamp) = 0;
  virtual void onHeartbeatFeed(const FeedTimer& timer) {}

 protected:
  Driver* driver_;
  ::nlohmann::json configs_;
  bool dry_run_;
  bool verbose_;
  std::string tag_;
  std::unordered_map<std::string, std::unique_ptr<LayeringExecutor3>> executors_;
  std::unordered_map<std::string, IOrderSubsystem*> os_;
};

class VmmAggArbitrageur : public VmmArbitrageur {
 public:
  VmmAggArbitrageur(
      const ::nlohmann::json& params,
      const ::nlohmann::json& configs,
      const bool& dry_run,
      const std::string& tag,
      const bool& verbose)
      : VmmArbitrageur(configs, dry_run, tag, verbose),
        thold_(params["thold_bp"].get<double>() * 1e-4),
        max_quote_(params["max_quote"]),
        order_freq_sec_(params["order_freq_sec"]),
        book_num_(params["book_num"]) {
    for (auto& config : configs_.items()) {
      std::string symbol = config.key();
      aggarb_on_[symbol] = config.value();
    }
  }

  bool onPricerUpdate(
      const std::unordered_map<std::string, VmmPriceResult>& result,
      const int64_t& ts) final {
    if (ts - last_order_ts_ < order_freq_sec_ * 1e9) return false;
    for (auto it1 = executors_.begin(); it1 != executors_.end(); ++it1) {
      const auto& symbol1 = it1->first;
      if (!aggarb_on_[symbol1]) continue;
      if (!os_[symbol1] || !os_[symbol1]->is_ready()) continue;
      if (result.find(symbol1) == result.end()) continue;
      for (auto it2 = std::next(it1); it2 != executors_.end(); ++it2) {
        const auto& symbol2 = it2->first;
        if (!aggarb_on_[symbol2]) continue;
        if (!os_[symbol2] || !os_[symbol2]->is_ready()) continue;
        if (result.find(symbol2) == result.end()) continue;
        const auto& bidp1 = result.find(symbol1)->second.native_bid.first;
        const auto& askp1 = result.find(symbol1)->second.native_ask.first;
        const auto& bidp2 = result.find(symbol2)->second.native_bid.first;
        const auto& askp2 = result.find(symbol2)->second.native_ask.first;
        if (!bidp1.has_value() || !bidp2.has_value() || !askp1.has_value() || !askp2.has_value())
          continue;
        std::string buy_symbol, sell_symbol;
        const BookBuilder *buy_book{nullptr}, *sell_book{nullptr};
        if (bidp2.value() / askp1.value() - 1.0 > thold_) {
          buy_symbol = symbol1;
          sell_symbol = symbol2;
          buy_book = result.find(symbol1)->second.bookbuilder;
          sell_book = result.find(symbol2)->second.bookbuilder;
        } else if (bidp1.value() / askp2.value() - 1.0 > thold_) {
          buy_symbol = symbol2;
          sell_symbol = symbol1;
          buy_book = result.find(symbol2)->second.bookbuilder;
          sell_book = result.find(symbol1)->second.bookbuilder;
        } else {
          continue;
        }
        const int maxN = std::min({book_num_,
                                   static_cast<int>(buy_book->AskSize()),
                                   static_cast<int>(sell_book->BidSize())});
        const auto& imSellp = InnermostOrderPrice(os_.at(buy_symbol)).second;
        const auto& imBuyp = InnermostOrderPrice(os_.at(sell_symbol)).first;
        int buyN = 0, sellN = 0;
        bool buyNfixed = false, sellNfixed = false;
        double avg_buyp = 0.0, avg_sellp = 0.0, buyp = 0.0, sellp = 0.0, buyq = 0.0, sellq = 0.0;
        while (buyN < maxN || sellN < maxN) {
          const auto& buy = buy_book->AskN(buyN);
          const auto& sell = sell_book->BidN(sellN);
          const double tmp_avg_buyp = (avg_buyp * buyq + buy->price * buy->qty) / (buyq + buy->qty);
          const double tmp_avg_sellp =
              (avg_sellp * sellq + sell->price * sell->qty) / (sellq + sell->qty);
          if (!buyNfixed && !sellNfixed && tmp_avg_sellp / tmp_avg_buyp - 1.0 > thold_ &&
              buy->price < imSellp && sell->price > imBuyp && avg_buyp * buyq < max_quote_ &&
              avg_sellp * sellq < max_quote_) {
            buyp = buy->price;
            sellp = sell->price;
            buyq += std::min(buy->qty, (max_quote_ - avg_buyp * buyq) / buyp);
            sellq += std::min(sell->qty, (max_quote_ - avg_sellp * sellq) / sellq);
            avg_buyp = tmp_avg_buyp;
            avg_sellp = tmp_avg_sellp;
            buyN++;
            sellN++;
          } else if (
              !sellNfixed && buyN > 0 && sellN > 0 &&
              tmp_avg_sellp / avg_buyp > avg_sellp / tmp_avg_buyp &&
              tmp_avg_sellp / avg_buyp - 1.0 > thold_ && sell->price > imBuyp &&
              avg_sellp * sellq < max_quote_ && buyq > sellq) {
            buyNfixed = true;
            avg_sellp = tmp_avg_sellp;
            sellp = sell->price;
            sellq += std::min(
                {sell->qty, (max_quote_ - avg_sellp * sellq) / sellp, std::max(buyq - sellq, 0.0)});
            sellN++;
          } else if (
              !buyNfixed && buyN > 0 && sellN > 0 &&
              avg_sellp / tmp_avg_buyp > tmp_avg_sellp / avg_buyp &&
              avg_sellp / tmp_avg_buyp - 1.0 > thold_ && buy->price < imSellp &&
              avg_buyp * buyq < max_quote_ && sellq > buyq) {
            sellNfixed = true;
            avg_buyp = tmp_avg_buyp;
            buyp = buy->price;
            buyq += std::min(
                {buy->qty, (max_quote_ - avg_buyp * buyq) / buyp, std::max(sellq - buyq, 0.0)});
            buyN++;
          } else {
            break;
          }
        }
        if (buyp == 0.0 || sellp == 0.0) continue;
        double buysmq = executors_[buy_symbol]->product_info().tick_qty().submittable_min_qty(buyp);
        double sellsmq =
            executors_[sell_symbol]->product_info().tick_qty().submittable_min_qty(sellp);
        const double smq = QTY_BUFFER * std::max(buysmq, sellsmq);
        const double qty = std::max(std::min(buyq, sellq), smq);
        const double spread = avg_sellp / avg_buyp - 1.0;
        SPDLOG_INFO("[AGGARB] {} {} {} {} {}", ts, spread, buyp, sellp, qty);
        if (!dry_run_) executors_[buy_symbol]->ManageAggOrders(ts, 1, buyp, qty, "AggArb");
        if (!dry_run_) executors_[sell_symbol]->ManageAggOrders(ts, -1, sellp, qty, "AggArb");
        last_order_ts_ = ts;
        return true;
      }
    }
    return false;
  }

 private:
  double thold_;
  double max_quote_;
  double order_freq_sec_;
  int book_num_;
  double QTY_BUFFER{1.2};
  int64_t last_order_ts_{0};
  std::unordered_map<std::string, bool> aggarb_on_;
};
}  // namespace coin2::strategy::vmm